import jwt
import time

from config.const import constant
from config.env.config_env import run_config
from forword.custom_error.token_errors import TokenExpireError


class TokenHandler(object):
    def __init__(self, num_max_token=-1, allow_expire=True, run_mode="dev"):
        self.jwt_salt = f'a plus inc#{run_mode}'
        self.encrypt_alg = 'HS256'
        self.exp_duration = constant.TOKEN_TIME_OUT
        self.headers = {
            "alg": "HS256",
            "typ": "JWT"
        }
        self.num_max_token = num_max_token
        self.allow_expire = allow_expire

    def gen_token(self, user_id):
        """Generate token"""
        payload = {constant.USER_ID_KEY: user_id}
        if self.allow_expire:
            exp_time = int(time.time() + self.exp_duration)
            payload.update({'exp': exp_time})

        token = jwt.encode(payload=payload, key=self.jwt_salt, algorithm=self.encrypt_alg, headers=self.headers)
        return token

    def parse_token(self, token):
        """Parse token"""
        try:
            info = jwt.decode(token, self.jwt_salt, algorithms=self.encrypt_alg)
        except Exception as e:
            raise TokenExpireError('请在header里Authorization字段中提供有效的登陆凭证')
        if info.get(constant.USER_ID_KEY) is None:
            raise TokenExpireError('请在header里Authorization字段中提供有效的登陆凭证')
        return info.get(constant.USER_ID_KEY)

    def update_token(self, payload):
        if self.allow_expire:
            if 'exp' in payload.keys():
                payload['exp'] += self.exp_duration
            else:
                payload['exp'] = time.time() + self.exp_duration
        token_new = jwt.encode(payload=payload, key=self.jwt_salt, algorithm=self.encrypt_alg, headers=self.headers)
        return token_new
